import ray

ray.init()
from process_pdf_enroll import BatchDocInfer

# Define an actor that uses 0.5 GPU
@ray.remote(num_gpus=1)
class GPUActor:
    def __init__(self):
        self.infer = BatchDocInfer("s3://llm-pdf-text-1/meta-index/scihub/v001/scihub/part-66210c190659-018706.jsonl", True, True)

    def compute(self):
        self.infer() 
        return 0

# Create actors
if 0:
    actor1 = GPUActor.remote()

    results = ray.get(actor1.compute.remote())
if 1:
    # Create actors
    actors = [GPUActor.remote() for _ in range(2)]

    # Run tasks on actors
    futures = [actor.compute.remote() for actor in actors]
    results = ray.get(futures)

print(results)


